"""
GPU资源管理API接口
"""
from typing import List, Optional, Dict, Any
from datetime import datetime
from fastapi import APIRouter, HTTPException, Depends, Query
import logging

from services.gpu_resource_manager import gpu_resource_manager
from core.deps import get_current_user
from schemas.user import User
from schemas.gpu_resource import (
    GPUInfo, GPUResourceRequest, GPUResourceAllocation,
    GPUClusterInfo, ResourceUsageStats, GPUHealthCheck,
    ResourceOptimizationSuggestion, LoadBalancingConfig
)

logger = logging.getLogger(__name__)
router = APIRouter()


@router.get("/info", response_model=Dict[int, GPUInfo])
async def get_gpu_info(
    gpu_id: Optional[int] = Query(None, description="特定GPU ID"),
    current_user: User = Depends(get_current_user)
):
    """获取GPU信息"""
    try:
        gpu_info = await gpu_resource_manager.get_gpu_info(gpu_id)
        return gpu_info
        
    except Exception as e:
        logger.error(f"获取GPU信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取GPU信息失败")


@router.get("/cluster", response_model=GPUClusterInfo)
async def get_cluster_info(
    current_user: User = Depends(get_current_user)
):
    """获取GPU集群信息"""
    try:
        cluster_info = await gpu_resource_manager.get_cluster_info()
        return cluster_info
        
    except Exception as e:
        logger.error(f"获取集群信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取集群信息失败")


@router.post("/allocate", response_model=GPUResourceAllocation)
async def allocate_resource(
    request: GPUResourceRequest,
    current_user: User = Depends(get_current_user)
):
    """分配GPU资源"""
    try:
        allocation = await gpu_resource_manager.allocate_resource(request)
        
        if not allocation:
            raise HTTPException(status_code=409, detail="无可用GPU资源")
        
        return allocation
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"分配GPU资源失败: {str(e)}")
        raise HTTPException(status_code=500, detail="分配GPU资源失败")


@router.delete("/allocate/{allocation_id}")
async def release_resource(
    allocation_id: str,
    current_user: User = Depends(get_current_user)
):
    """释放GPU资源"""
    try:
        success = await gpu_resource_manager.release_resource(allocation_id)
        
        if not success:
            raise HTTPException(status_code=404, detail="分配记录不存在")
        
        return {"message": "资源释放成功", "allocation_id": allocation_id}
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"释放GPU资源失败: {str(e)}")
        raise HTTPException(status_code=500, detail="释放GPU资源失败")


@router.get("/usage/{gpu_id}", response_model=List[ResourceUsageStats])
async def get_usage_stats(
    gpu_id: int,
    hours: int = Query(24, ge=1, le=168, description="统计时间范围(小时)"),
    current_user: User = Depends(get_current_user)
):
    """获取GPU使用统计"""
    try:
        stats = await gpu_resource_manager.get_usage_stats(gpu_id, hours)
        return stats
        
    except Exception as e:
        logger.error(f"获取使用统计失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取使用统计失败")


@router.post("/health-check", response_model=Dict[int, GPUHealthCheck])
async def perform_health_check(
    gpu_id: Optional[int] = Query(None, description="特定GPU ID"),
    current_user: User = Depends(get_current_user)
):
    """执行GPU健康检查"""
    try:
        health_checks = await gpu_resource_manager.perform_health_check(gpu_id)
        return health_checks
        
    except Exception as e:
        logger.error(f"健康检查失败: {str(e)}")
        raise HTTPException(status_code=500, detail="健康检查失败")


@router.get("/optimization-suggestions", response_model=List[ResourceOptimizationSuggestion])
async def get_optimization_suggestions(
    current_user: User = Depends(get_current_user)
):
    """获取资源优化建议"""
    try:
        suggestions = await gpu_resource_manager.get_optimization_suggestions()
        return suggestions
        
    except Exception as e:
        logger.error(f"获取优化建议失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取优化建议失败")


@router.get("/load-balancing-config", response_model=LoadBalancingConfig)
async def get_load_balancing_config(
    current_user: User = Depends(get_current_user)
):
    """获取负载均衡配置"""
    try:
        config = gpu_resource_manager.load_balancing_config
        return config
        
    except Exception as e:
        logger.error(f"获取负载均衡配置失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取负载均衡配置失败")


@router.put("/load-balancing-config")
async def update_load_balancing_config(
    config: LoadBalancingConfig,
    current_user: User = Depends(get_current_user)
):
    """更新负载均衡配置"""
    try:
        gpu_resource_manager.load_balancing_config = config
        return {"message": "负载均衡配置更新成功"}
        
    except Exception as e:
        logger.error(f"更新负载均衡配置失败: {str(e)}")
        raise HTTPException(status_code=500, detail="更新负载均衡配置失败")


@router.post("/detect")
async def detect_gpus(
    current_user: User = Depends(get_current_user)
):
    """重新检测GPU设备"""
    try:
        await gpu_resource_manager.detect_gpus()
        gpu_info = await gpu_resource_manager.get_gpu_info()
        
        return {
            "message": "GPU检测完成",
            "gpu_count": len(gpu_info),
            "gpus": gpu_info
        }
        
    except Exception as e:
        logger.error(f"GPU检测失败: {str(e)}")
        raise HTTPException(status_code=500, detail="GPU检测失败")


@router.get("/allocations")
async def get_current_allocations(
    current_user: User = Depends(get_current_user)
):
    """获取当前资源分配情况"""
    try:
        allocations = {}
        for allocation_id, allocation in gpu_resource_manager.allocations.items():
            allocations[allocation_id] = {
                "allocation_id": allocation.allocation_id,
                "gpu_id": allocation.gpu_id,
                "memory_allocated": allocation.memory_allocated,
                "task_id": allocation.task_id,
                "allocated_at": allocation.allocated_at,
                "expires_at": allocation.expires_at,
                "exclusive": allocation.exclusive
            }
        
        return {
            "total_allocations": len(allocations),
            "allocations": allocations
        }
        
    except Exception as e:
        logger.error(f"获取分配情况失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取分配情况失败")


@router.get("/metrics")
async def get_gpu_metrics(
    current_user: User = Depends(get_current_user)
):
    """获取GPU性能指标"""
    try:
        cluster_info = await gpu_resource_manager.get_cluster_info()
        gpu_info = await gpu_resource_manager.get_gpu_info()
        
        # 计算各种指标
        metrics = {
            "cluster": {
                "total_gpus": cluster_info.total_gpus,
                "available_gpus": cluster_info.available_gpus,
                "busy_gpus": cluster_info.busy_gpus,
                "error_gpus": cluster_info.error_gpus,
                "availability_rate": cluster_info.available_gpus / max(cluster_info.total_gpus, 1) * 100,
                "average_utilization": cluster_info.average_utilization,
                "memory_usage_rate": cluster_info.used_memory / max(cluster_info.total_memory, 1) * 100
            },
            "gpus": {}
        }
        
        for gpu_id, gpu in gpu_info.items():
            metrics["gpus"][gpu_id] = {
                "name": gpu.name,
                "utilization": gpu.utilization,
                "memory_usage_rate": gpu.memory_used / max(gpu.memory_total, 1) * 100,
                "temperature": gpu.temperature,
                "power_usage": gpu.power_usage,
                "status": gpu.status
            }
        
        return metrics
        
    except Exception as e:
        logger.error(f"获取GPU指标失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取GPU指标失败")


@router.post("/benchmark")
async def run_gpu_benchmark(
    gpu_id: Optional[int] = Query(None, description="特定GPU ID"),
    duration: int = Query(60, ge=10, le=300, description="基准测试时长(秒)"),
    current_user: User = Depends(get_current_user)
):
    """运行GPU基准测试"""
    try:
        # 这里应该实现实际的基准测试逻辑
        # 目前返回模拟数据
        
        gpu_ids = [gpu_id] if gpu_id is not None else list(gpu_resource_manager.gpus.keys())
        results = {}
        
        for gid in gpu_ids:
            if gid not in gpu_resource_manager.gpus:
                continue
            
            # 模拟基准测试结果
            results[gid] = {
                "gpu_id": gid,
                "compute_score": 8500 + gid * 100,  # 模拟计算分数
                "memory_bandwidth": 900 + gid * 50,  # 模拟内存带宽 GB/s
                "fp32_performance": 25.5 + gid * 2.5,  # 模拟FP32性能 TFLOPS
                "fp16_performance": 51.0 + gid * 5.0,  # 模拟FP16性能 TFLOPS
                "duration": duration,
                "timestamp": datetime.now()
            }
        
        return {
            "message": "基准测试完成",
            "results": results
        }
        
    except Exception as e:
        logger.error(f"GPU基准测试失败: {str(e)}")
        raise HTTPException(status_code=500, detail="GPU基准测试失败")


@router.post("/maintenance/{gpu_id}")
async def set_gpu_maintenance(
    gpu_id: int,
    maintenance: bool = Query(..., description="是否进入维护模式"),
    current_user: User = Depends(get_current_user)
):
    """设置GPU维护模式"""
    try:
        if gpu_id not in gpu_resource_manager.gpus:
            raise HTTPException(status_code=404, detail="GPU不存在")
        
        from schemas.gpu_resource import GPUStatus
        
        if maintenance:
            gpu_resource_manager.gpus[gpu_id].status = GPUStatus.MAINTENANCE
            message = f"GPU {gpu_id} 已进入维护模式"
        else:
            gpu_resource_manager.gpus[gpu_id].status = GPUStatus.AVAILABLE
            message = f"GPU {gpu_id} 已退出维护模式"
        
        return {"message": message}
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"设置GPU维护模式失败: {str(e)}")
        raise HTTPException(status_code=500, detail="设置GPU维护模式失败")


@router.get("/temperature-history/{gpu_id}")
async def get_temperature_history(
    gpu_id: int,
    hours: int = Query(24, ge=1, le=168, description="历史时间范围(小时)"),
    current_user: User = Depends(get_current_user)
):
    """获取GPU温度历史"""
    try:
        stats = await gpu_resource_manager.get_usage_stats(gpu_id, hours)
        
        temperature_data = [
            {
                "timestamp": stat.timestamp,
                "temperature": stat.temperature
            }
            for stat in stats if stat.temperature is not None
        ]
        
        return {
            "gpu_id": gpu_id,
            "data_points": len(temperature_data),
            "temperature_history": temperature_data
        }
        
    except Exception as e:
        logger.error(f"获取温度历史失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取温度历史失败")


@router.get("/power-history/{gpu_id}")
async def get_power_history(
    gpu_id: int,
    hours: int = Query(24, ge=1, le=168, description="历史时间范围(小时)"),
    current_user: User = Depends(get_current_user)
):
    """获取GPU功耗历史"""
    try:
        stats = await gpu_resource_manager.get_usage_stats(gpu_id, hours)
        
        power_data = [
            {
                "timestamp": stat.timestamp,
                "power_usage": stat.power_usage
            }
            for stat in stats if stat.power_usage is not None
        ]
        
        return {
            "gpu_id": gpu_id,
            "data_points": len(power_data),
            "power_history": power_data
        }
        
    except Exception as e:
        logger.error(f"获取功耗历史失败: {str(e)}")
        raise HTTPException(status_code=500, detail="获取功耗历史失败")